## Event study plot function
es_plot <- function(es_results){
  g <- tidy(es_results)
  coefs <- c(es_results$att.egt)
  ses <- c(es_results$se.egt)
  months <- c(es_results$egt)
  coefs.df <- as.data.frame(cbind(coefs, ses, months))
  coefs.df$Timing <- ifelse(coefs.df$months>=0, "After","Before")
  ggplot(coefs.df, aes(months,coefs,col=Timing)) + 
    geom_hline(yintercept=0, lty=2, lwd=.5, colour="grey25") +
    geom_point(size=3) +
    geom_errorbar(aes(ymin=coefs - 1.96*ses, ymax=coefs + 1.96*ses), 
                  lwd=.5, width=0) +
    theme_minimal() +
    scale_color_manual(values=c("firebrick4","#999999")) +
    labs(x="Months to approval",y="Change in registration") +
    scale_x_continuous(limits = c(-12,12),
                       breaks = seq(-12,12,1)) +
    scale_y_continuous(breaks = function(z) seq(-0.02, range(z)[2], by = 0.02))
}

es_plot_pre <- function(es_results){
  g <- tidy(es_results)
  coefs <- c(es_results$att.egt)
  ses <- c(es_results$se.egt)
  months <- c(es_results$egt)
  coefs.df <- as.data.frame(cbind(coefs, ses, months))
  coefs.df$Timing <- ifelse(coefs.df$months>=0, "After","Before")
  ggplot(coefs.df, aes(months,coefs,col=Timing)) + 
    geom_hline(yintercept=0, lty=2, lwd=.5, colour="grey25") +
    geom_point(size=3) +
    geom_errorbar(aes(ymin=coefs - 1.96*ses, ymax=coefs + 1.96*ses), 
                  lwd=.5, width=0) +
    theme_minimal() +
    scale_color_manual(values=c("firebrick4","#999999")) +
    labs(x="Months to approval",y="Change in registration") +
    scale_x_continuous(limits = c(-48,12),
                       breaks = seq(-48,12,1)) +
    scale_y_continuous(breaks = function(z) seq(-0.2, range(z)[2], by = 0.1))
}

es_plot_post <- function(es_results){
  g <- tidy(es_results)
  coefs <- c(es_results$att.egt)
  ses <- c(es_results$se.egt)
  months <- c(es_results$egt)
  coefs.df <- as.data.frame(cbind(coefs, ses, months))
  coefs.df$Timing <- ifelse(coefs.df$months>=0, "After","Before")
  ggplot(coefs.df, aes(months,coefs,col=Timing)) + 
    geom_hline(yintercept=0, lty=2, lwd=.5, colour="grey25") +
    geom_point(size=3) +
    geom_errorbar(aes(ymin=coefs - 1.96*ses, ymax=coefs + 1.96*ses), 
                  lwd=.5, width=0) +
    theme_minimal() +
    scale_color_manual(values=c("firebrick4","#999999")) +
    labs(x="Months to approval",y="Change in registration") +
    scale_x_continuous(limits = c(-12,48),
                       breaks = seq(-12,48,1)) +
    scale_y_continuous(breaks = function(z) seq(-0.2, range(z)[2], by = 0.1))
}
## Honest DID Function (from https://github.com/asheshrambachan/HonestDiD)
#' @title honest_did
#'
#' @description a function to compute a sensitivity analysis
#'  using the approach of Rambachan and Roth (2021)
honest_did <- function(...) UseMethod("honest_did")

#' @title honest_did.AGGTEobj
#'
#' @description a function to compute a sensitivity analysis
#'  using the approach of Rambachan and Roth (2021) when
#'  the event study is estimating using the `did` package
#'
#' @param e event time to compute the sensitivity analysis for.
#'  The default value is `e=0` corresponding to the "on impact"
#'  effect of participating in the treatment.
#' @param type Options are "smoothness" (which conducts a
#'  sensitivity analysis allowing for violations of linear trends
#'  in pre-treatment periods) or "relative_magnitude" (which
#'  conducts a sensitivity analysis based on the relative magnitudes
#'  of deviations from parallel trends in pre-treatment periods).
#' @inheritParams HonestDiD::createSensitivityResults
#' @inheritParams HonestDid::createSensitivityResults_relativeMagnitudes
honest_did.AGGTEobj <- function(es,
                                e          = 0,
                                type       = c("smoothness", "relative_magnitude"),
                                gridPoints = 100,
                                ...) {
  
  type <- match.arg(type)
  
  # Make sure that user is passing in an event study
  if (es$type != "dynamic") {
    stop("need to pass in an event study")
  }
  
  # Check if used universal base period and warn otherwise
  if (es$DIDparams$base_period != "universal") {
    stop("Use a universal base period for honest_did")
  }
  
  # Recover influence function for event study estimates
  es_inf_func <- es$inf.function$dynamic.inf.func.e
  
  # Recover variance-covariance matrix
  n <- nrow(es_inf_func)
  V <- t(es_inf_func) %*% es_inf_func / n / n
  
  # Remove the coefficient normalized to zero
  referencePeriodIndex <- which(es$egt == -1)
  V    <- V[-referencePeriodIndex,-referencePeriodIndex]
  beta <- es$att.egt[-referencePeriodIndex]
  
  nperiods <- nrow(V)
  npre     <- sum(1*(es$egt < -1))
  npost    <- nperiods - npre
  baseVec1 <- basisVector(index=(e+1),size=npost)
  orig_ci  <- constructOriginalCS(betahat        = beta,
                                  sigma          = V,
                                  numPrePeriods  = npre,
                                  numPostPeriods = npost,
                                  l_vec          = baseVec1)
  
  if (type=="relative_magnitude") {
    robust_ci <- createSensitivityResults_relativeMagnitudes(betahat        = beta,
                                                             sigma          = V,
                                                             numPrePeriods  = npre,
                                                             numPostPeriods = npost,
                                                             l_vec          = baseVec1,
                                                             gridPoints     = gridPoints,
                                                             ...)
    
  } else if (type == "smoothness") {
    robust_ci <- createSensitivityResults(betahat        = beta,
                                          sigma          = V,
                                          numPrePeriods  = npre,
                                          numPostPeriods = npost,
                                          l_vec          = baseVec1,
                                          ...)
  }
  
  return(list(robust_ci=robust_ci, orig_ci=orig_ci, type=type))
}

# for creating tables
extractFromDataFrame <- function (dataFrame) {
  coef.row.indices <- seq(1, nrow(dataFrame), 1)
  se.row.indices <- seq(1, nrow(dataFrame), 1)
  texregObjects <- list()
  coefs <- dataFrame[coef.row.indices, 4]
  coefnames <- as.character(dataFrame[coef.row.indices, 3])
  se <- dataFrame[se.row.indices, 5]
  tr <- createTexreg(coef = coefs, coef.names = coefnames, se = se)
  texregObjects <- list(tr)
  return(texregObjects)
}

make_coefs_dfs <- function(results, label) {
  coefs <- c(results$att.egt)
  ses <- c(results$se.egt)
  months <- c(results$egt)
  hte <- label
  coefs <- as.data.frame(cbind(coefs, ses, months, hte)) 
  coefs$Timing <- ifelse(coefs$months>=0, "After","Before")
  coefs$ses <- as.numeric(coefs$ses)
  coefs$coefs <- as.numeric(coefs$coefs)
  coefs$months <- as.numeric(coefs$months)
  return(coefs)
}

f <- function(x) factor(x, levels = unique(x))

make_coef_plot <- function(coefs1, coefs2, coefs3) {
  if(missing(coefs3)) {
    df <- bind_rows(coefs1, coefs2)
    ggplot(df,
           aes(x = months, y = coefs, color = f(hte), shape=f(hte))) +
      geom_hline(yintercept=0, lwd=.6, colour="grey50") +
      geom_point(size=3,position=position_dodge(width=1)) +
      geom_errorbar(aes(ymin=coefs - 1.96*ses, ymax=coefs + 1.96*ses), 
                    lwd=.6, width=0,position=position_dodge(width=1), inherit.aes = FALSE) +
      theme_minimal() +
      gghighlight(Timing == "After",use_direct_label = FALSE,
                  use_group_by = FALSE) +
      scale_color_manual(values=c("purple","black", "#69b3a2"))+
      labs(x="Months to approval",y="Change in registration", color=NULL,shape=NULL) +
      scale_x_continuous(limits = c(-12.6,12.6),
                         breaks = seq(-12,12,1)) +
      scale_y_continuous(breaks = function(z) seq(-0.02, range(z)[2], by = 0.02))
  } else {
    df <- bind_rows(coefs1, coefs2, coefs3)
    ggplot(df,
           aes(x = months, y = coefs, color = f(hte), shape=f(hte))) +
      geom_hline(yintercept=0, lwd=.6, colour="grey50") +
      geom_point(size=3,position=position_dodge(width=1)) +
      geom_errorbar(aes(ymin=coefs - 1.96*ses, ymax=coefs + 1.96*ses), 
                    lwd=.6, width=0,position=position_dodge(width=1), inherit.aes = FALSE) +
      theme_minimal() +
      gghighlight(Timing == "After",use_direct_label = FALSE,
                  use_group_by = FALSE) +
      scale_color_manual(values=c("purple","black","#69b3a2"))+
      labs(x="Months to approval",y="Change in registration", color=NULL,shape=NULL) +
      scale_x_continuous(limits = c(-12.6,12.6),
                         breaks = seq(-12,12,1)) +
      scale_y_continuous(breaks = function(z) seq(-0.02, range(z)[2], by = 0.02))
  }
 }
